import json
import uuid

from fastapi import FastAPI, WebSocket, WebSocketDisconnect

from .resources import ResourceManager
from .schemas import MCPRequest, MCPResponse, ToolInvocation
from .tools import tool_manager

app = FastAPI(title="Python MCP Server")
resource_manager = ResourceManager()


# 示例工具注册
@tool_manager.register_tool
async def get_weather(city: str) -> str:
    """获取指定城市的天气信息"""
    # 这里应该是调用天气API的实现
    return f"Weather in {city}: Sunny, 25°C"


from fastapi import Depends, HTTPException, status
from fastapi.security import HTTPBasic, HTTPBasicCredentials

security = HTTPBasic()


def authenticate(credentials: HTTPBasicCredentials = Depends(security)):
    correct_username = "admin"
    correct_password = "secret"
    if (credentials.username != correct_username or
            credentials.password != correct_password):
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Incorrect credentials",
            headers={"WWW-Authenticate": "Basic"},
        )
    return credentials.username


# WebSocket端点
@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket, token: str = None):
    if token != "my-secret-token":
        await websocket.close(code=status.WS_1008_POLICY_VIOLATION)
        return
    await websocket.accept()
    try:
        while True:
            data = await websocket.receive_text()
            try:
                request = MCPRequest(**json.loads(data))

                if request.action == "list_tools":
                    response = MCPResponse(
                        request_id=request.request_id,
                        success=True,
                        data=tool_manager.list_tools()
                    )
                elif request.action == "invoke_tool":
                    invocation = ToolInvocation(**request.payload)
                    result = await tool_manager.invoke_tool(
                        invocation.tool_name,
                        invocation.parameters
                    )
                    response = MCPResponse(
                        request_id=request.request_id,
                        success=result.success,
                        data=result.output if result.success else None,
                        error=result.error
                    )
                else:
                    response = MCPResponse(
                        request_id=request.request_id,
                        success=False,
                        error=f"Unknown action: {request.action}"
                    )

                await websocket.send_text(response.json())
            except Exception as e:
                error_response = MCPResponse(
                    request_id=getattr(request, 'request_id', str(uuid.uuid4())),
                    success=False,
                    error=str(e)
                )
                await websocket.send_text(error_response.json())

    except WebSocketDisconnect:
        print("Client disconnected")


# HTTP端点
@app.get("/tools")
async def list_tools():
    return tool_manager.list_tools()


@app.post("/invoke")
async def invoke_tool(invocation: ToolInvocation):
    result = await tool_manager.invoke_tool(invocation.tool_name, invocation.parameters)
    return result


if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="0.0.0.0", port=8000)
